import torch
import torch.nn.functional as F
from utils import Adder,class_indices,write_and_print_logs
from dice_score import multiclass_iou
from dist import is_primary,barrier
def _valid(model, args, ep, data_loader_val):
    barrier()
    model.eval()
    if is_primary():
        firing_rate_mean_adder = Adder()
        TP_percent_mean_adder = Adder()
        FP_2_TP_mean_adder = Adder()
        iou_cls_adder = [Adder() for _ in range(args.num_classes)]
        TP_percent_cls_adder = [Adder() for _ in range(args.num_classes)]
        FP_2_TP_cls_adder = [Adder() for _ in range(args.num_classes)]
        iou_mean_adder = Adder()
    with torch.no_grad():
        for idx, data in enumerate(data_loader_val):
            input_img, label_img = data
            input_img = input_img.to(dtype=torch.float32).cuda()
            label_img = label_img.to(dtype=torch.float32).cuda()
            label_img = label_img.squeeze(1)
            label_img = torch.tensor(label_img.cpu().numpy(),dtype=torch.long).cuda()
            label_img = F.one_hot(label_img, args.num_classes).permute(0, 3, 1, 2).float()

            pred,spike_mat = model(input_img,True)
            pred = pred[2]
            pred = F.one_hot(pred.argmax(dim=1), args.num_classes).permute(0, 3, 1, 2).float()

            iou_score_list,TP_percent_list,FP_2_TP_list = multiclass_iou(pred, label_img, reduce_batch_first=True)
            if is_primary():
                firing_rate_mean_adder(spike_mat)
                iou_temp_adder = Adder()
                for index,i in enumerate(iou_score_list):
                    iou_cls_adder[index](i)
                    iou_temp_adder(i)
                TP_percent_temp_adder = Adder()
                for index,i in enumerate(TP_percent_list):
                    TP_percent_cls_adder[index](i)
                    TP_percent_temp_adder(i)
                FP_2_TP_temp_adder = Adder()
                for index,i in enumerate(FP_2_TP_list):
                    FP_2_TP_cls_adder[index](i)
                    FP_2_TP_temp_adder(i)
                iou_mean_adder(iou_temp_adder.average())
                TP_percent_mean_adder(TP_percent_temp_adder.average())
                FP_2_TP_mean_adder(FP_2_TP_temp_adder.average())

    if is_primary():
        write_and_print_logs(args.text_logs_dir,"ep: %d, mean firing rate: %f"%(ep,firing_rate_mean_adder.average()))
        for key,value in class_indices.items():
            write_and_print_logs(args.text_logs_dir,'%s iou: %.3f TP_per: %.3f FP_2_TP: %.3f'%(\
                key,iou_cls_adder[value].average(),TP_percent_cls_adder[value].average(),FP_2_TP_cls_adder[value].average()))
    model.train()
    barrier()
    if is_primary():
        return iou_mean_adder.average()
    else:
        return 0
